#include "BreakKey.h"

PyObject* pyf_SetV(PyObject* self, PyObject* args);
PyObject* pyf_GetV(PyObject* self, PyObject* args);
PyObject* pyf_SetPlainText(PyObject* self, PyObject* args);
PyObject* pyf_GetPlainText(PyObject* self, PyObject* args);
PyObject* pyf_AddAes(PyObject* self, PyObject* args);
PyObject* pyf_debugout(PyObject* self, PyObject* args);
static PyMethodDef BreakKeyMethods[] = {
	{"SetV",			pyf_SetV,			METH_VARARGS, ""},
	{"GetV",			pyf_GetV,			METH_VARARGS, ""},
	{"SetPlainText",	pyf_SetPlainText,	METH_VARARGS, ""},
	{"GetPlainText",	pyf_GetPlainText,	METH_VARARGS, ""},
	{"AddAes",			pyf_AddAes,			METH_VARARGS, ""},
	{"debugoutstring",	pyf_debugout,		METH_VARARGS, ""},
	{NULL, NULL, 0, NULL}       
};

BreakKey* BreakKey::msInstance = NULL;

BreakKey::BreakKey():mCurrCount(0),mNodeCount(0),mAesCount(0),mFreeNodeCount(0),
	mTopNode(NULL),mBotNode(NULL),mCurrNode(NULL),
	mFreeTopNode(NULL),mFreeBotNode(NULL)
{
}

void BreakKey::Init()
{
	InitializeCriticalSection(&mLock);
	InitializeCriticalSection(&mLockCurr);
	InitializeCriticalSection(&mLockFree);
	InitializeCriticalSection(&mLockDebugOut);
	InitializeCriticalSection(&mLockCmd);
	memset(mAes, 0, sizeof(mAes));
	
	Py_Initialize();
	Py_InitModule("BreakKey", BreakKeyMethods);
	ExecuteScript("init");	
	mBatchCount = GetV(1) * GetV(2);
	mMaxBatch = GetV(3);
	mAABB = CreateAABB();
	mAABB->SetHost(this);
	mAABB->Start();	
	PushCmd("from BreakKey import *");
	PushCmd("def debugout(msg):\r\n   debugoutstring(str(msg))\r\n");
	mLastNewAABBTick = GetTickCount();
	mLastRunTick = GetTickCount();
}

void BreakKey::Uninit()
{
	Py_Finalize();
	DeleteCriticalSection(&mLockCmd);
	DeleteCriticalSection(&mLockDebugOut);
	DeleteCriticalSection(&mLockFree);
	DeleteCriticalSection(&mLockCurr);
	DeleteCriticalSection(&mLock);
}

void BreakKey::Run()
{
	mLastRunTick = GetTickCount();	
	EnterCriticalSection(&mLockCmd);
	if (mCmds.size() > 0)
	{
		for (std::vector<std::string>::iterator itr = mCmds.begin();
			itr != mCmds.end(); itr++)
		{
			PyRun_SimpleString(itr->c_str());
		}
		mCmds.clear();
	}
	LeaveCriticalSection(&mLockCmd);
}

bool BreakKey::ExecuteScript(const char* szScriptFile)
{
	if (!PyImport_ImportModule(szScriptFile))
	{		
		DebugOut("ExecuteScript %s error ", szScriptFile);
		PyErr_Print();
		DebugOut("\n");
	}	
	return true;
}

void BreakKey::SetV(int iKey, int iValue)
{
	mKeyValues[iKey] = iValue;
}

int BreakKey::GetV(int iKey)
{
	if (mKeyValues.find(iKey) != mKeyValues.end())
	{
		return mKeyValues[iKey];
	}
	return -1;
}

void BreakKey::SetPlainText(const std::string& strPlainText)
{
	mPlainText = strPlainText;
}

std::string BreakKey::GetPlainText()
{
	return mPlainText;
}

BlockNode* BreakKey::PopBlockNode()
{
	BlockNode* node = NULL;
	EnterCriticalSection(&mLock);
	if (mTopNode)
	{
		node = mTopNode;
		mTopNode = mTopNode->mNext;
		if (!mTopNode)
		{
			mBotNode = NULL;
		}
		mNodeCount--;
	}
	LeaveCriticalSection(&mLock);
	return node;
}

void BreakKey::FreeBlockNode(BlockNode* node)
{
	EnterCriticalSection(&mLockFree);
	if (mFreeBotNode)
	{
		mFreeBotNode->mNext = node;
		mFreeBotNode = node;
		node->mNext = NULL;
	}else
	{
		mFreeTopNode = node;
		mFreeBotNode = node;
		node->mNext = NULL;
	}
	mFreeNodeCount++;
	LeaveCriticalSection(&mLockFree);
}

IAABB* BreakKey::AABB()
{
	return mAABB;
}

BlockNode* BreakKey::GetBlockNode()
{
	BlockNode* node = NULL;
	EnterCriticalSection(&mLockFree);
	if (mFreeTopNode)
	{
		node = mFreeTopNode;
		mFreeTopNode = mFreeTopNode->mNext;		
		if (!mFreeTopNode)
		{
			mFreeBotNode = NULL;
		}
		mFreeNodeCount--;
	}
	LeaveCriticalSection(&mLockFree);
	if (!node)
	{
		while(mNodeCount>10)
			Sleep(10);
		node = new BlockNode;
		node->mData = new AES_KEYTYPE16[mBatchCount];		
		node->mNext = NULL;
	}
	memset(node->mData, 0, sizeof(AES_KEYTYPE16) * mBatchCount);
	return node;
}

void __stdcall BreakKey::NewAABB(const char* szAABB, int nLen, int nCount)
{
	mLastNewAABBTick = GetTickCount();
	EnterCriticalSection(&mLockCurr);
	if (!mCurrNode)
		mCurrNode = GetBlockNode();
	for (int i = 0; i < nCount; i++)
	{
		for(int k = 0;k < nLen; k++)
			mCurrNode->mData[mCurrCount].keys[k] = szAABB[i*nLen + k];
		for(int k = nLen;k < 16; k++)
			mCurrNode->mData[mCurrCount].keys[k] = 0;	
		mCurrCount++;
		if (mCurrCount >= mBatchCount)
		{
			mCurrNode->mNext = NULL;
			EnterCriticalSection(&mLock);
			if (mBotNode)
			{
				mBotNode->mNext = mCurrNode;
				mBotNode = mCurrNode;
			}else
			{
				mTopNode = mCurrNode;
				mBotNode = mCurrNode;				
			}
			mNodeCount++;
			LeaveCriticalSection(&mLock);
			mCurrCount = 0;
			mCurrNode = GetBlockNode();
		}
	}		
	LeaveCriticalSection(&mLockCurr);
}

void BreakKey::AddAes(const char* szAes)
{
	if (strlen(szAes) != 32)
	{
		DebugOut("[error]AddAes len error szAes=%s len=%d", szAes, strlen(szAes));
		return;
	}

	int h, l;
	for (int i = 0; i < 32/2; i++)
	{		
		switch(*(szAes + i*2))
		{
		case '0': case '1': case '2': case '3': case '4': 
		case '5': case '6': case '7': case '8': case '9':
			h = *(szAes + i*2) - '0';
			break;
		case 'a': case 'b': case 'c': case 'd': case 'e': case 'f':
			h = 10 + *(szAes + i*2) - 'a';
			break;
		case 'A': case 'B': case 'C': case 'D': case 'E': case 'F':
			h = 10 + *(szAes + i*2) - 'A';
			break;
		default:
			DebugOut("[error]AddAes error aes format szAes=%s pos=%d", szAes, i*2);
			return;
		}	

		switch(*(szAes + i*2 + 1))
		{
		case '0': case '1': case '2': case '3': case '4': 
		case '5': case '6': case '7': case '8': case '9':
			l = *(szAes + i*2 + 1) - '0';
			break;
		case 'a': case 'b': case 'c': case 'd': case 'e': case 'f':
			l = 10 + *(szAes + i*2 + 1) - 'a';
			break;
		case 'A': case 'B': case 'C': case 'D': case 'E': case 'F':
			l = 10 + *(szAes + i*2 + 1) - 'A';
			break;
		default:
			DebugOut("[error]AddAes error aes format szAes=%s pos=%d", szAes, i*2+1);
			return;
		}

		mAes[mAesCount].keys[i] = h*16 + l;
	}
	mAesCount++;	
}

void BreakKey::DebugOut(const char * _Format, ...)
{
	va_list ap;
	va_start(ap, _Format);	
	char buf[512];
	int iRet = vsnprintf(buf, 512, _Format, ap);
	if (iRet < 0)
	{
		sprintf(buf, "DebugOut Error iRet=%d", iRet);
	}
	va_end(ap);

	EnterCriticalSection(&mLockDebugOut);
	mDebugMsgs.push_back(buf);
	LeaveCriticalSection(&mLockDebugOut);
}

void BreakKey::RunDebugOut()
{
	EnterCriticalSection(&mLockDebugOut);
	if (mDebugMsgs.size() > 0)
	{
		for (std::vector<std::string>::iterator itr = mDebugMsgs.begin();
			itr != mDebugMsgs.end(); itr++)
		{
			OutputDebugString(itr->c_str());
		}
		mDebugMsgs.clear();
	}
	LeaveCriticalSection(&mLockDebugOut);
}

void BreakKey::PushCmd(const char *szCmd)
{
	if (strcmp(szCmd, "info") == 0)
	{
		__int64 curNum = 0, totalNum = 0;
		mAABB->GetProccess(&curNum, &totalNum);
		printf("Progress:%lld/%lld BlockNum:%lu FreeBlockNum:%lu AABBTick:%lu RunTick:%lu\n", 
			curNum, totalNum, mNodeCount, mFreeNodeCount, mLastNewAABBTick, mLastRunTick);
	}else
	{
		EnterCriticalSection(&mLockCmd);
		mCmds.push_back(szCmd);
		LeaveCriticalSection(&mLockCmd);
	}	
}

PyObject* pyf_SetV(PyObject* self, PyObject* args) 
{
	int iKey, iValue;

	iValue = iKey = 0;
	if (! PyArg_ParseTuple(args, "ii", &iKey, &iValue))
		return NULL;

	BreakKey::Instance()->SetV(iKey, iValue);
	Py_IncRef(Py_None);
	return Py_None;
}

PyObject* pyf_GetV(PyObject* self, PyObject* args) 
{
	int iKey, iResult;

	iKey = 0;
	if (! PyArg_ParseTuple(args, "i", &iKey))
		return NULL;

	iResult = BreakKey::Instance()->GetV(iKey);	
	return Py_BuildValue("i", iResult);
}

PyObject* pyf_SetPlainText(PyObject* self, PyObject* args)
{
	char* szText = NULL;
	if (! PyArg_ParseTuple(args, "s", &szText))
		return NULL;

	BreakKey::Instance()->SetPlainText(szText);
	Py_IncRef(Py_None);
	return Py_None;
}

PyObject* pyf_GetPlainText(PyObject* self, PyObject* args)
{
	std::string strText = BreakKey::Instance()->GetPlainText(); 
	return Py_BuildValue("s", strText.c_str());
}

PyObject* pyf_AddAes(PyObject* self, PyObject* args)
{
	char* szText = NULL;
	if (! PyArg_ParseTuple(args, "s", &szText))
		return NULL;

	BreakKey::Instance()->AddAes(szText);
	Py_IncRef(Py_None);
	return Py_None;
}

PyObject* pyf_debugout(PyObject* self, PyObject* args)
{
	char* szText = NULL;
	if (! PyArg_ParseTuple(args, "s", &szText))
		return NULL;

	BreakKey::Instance()->DebugOut("%s", szText);
	Py_IncRef(Py_None);
	return Py_None;
}
